import torch
import wandb
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint,Callback
from config import wandb_logger

class LogPredictionSamplesCallback(Callback):
    def on_validation_batch_end(
            self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
        """Called when the validation batch ends."""
        # `outputs` comes from `LightningModule.validation_step`
        # which corresponds to our model predictions in this case
        # Let's log 20 sample image predictions from the first batch
        if batch_idx == 0:
            n = 4
            x, y = batch
            # # Option 1: log images with `WandbLogger.log_image`
            images = [img for img in x[:n]]
            captions = [f'Ground Truth: {y_i} - Prediction: {y_pred}'
                        for y_i, y_pred in zip(y[:n], outputs["preds"][:n])]
            wandb_logger.log_image(
                key='sample_images',
                images=images,
                caption=captions
            )
            # Option 2: log images and predictions as a W&B Table
            columns = ['image', 'ground truth', 'prediction']
            data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs["preds"][:n]))]

            wandb_logger.log_table(
                key='sample_table',
                columns=columns,
                data=data)

            # data should be a list of lists
            columns = ["input", "label", "prediction"]
            my_data = [["cheese", "english", "english"], ["fromage", "french", "spanish"]]

            # using columns and data
            wandb_logger.log_text(key="my_samples", columns=columns, data=my_data)
            import pandas as pd
            my_dataframe = pd.DataFrame(my_data)
            # using a pandas DataFrame
            wandb_logger.log_text(key="my_samples", dataframe=my_dataframe)